`define SIZE_TAG   21:0
`define SIZE_INDEX 7:0
`define SIZE_BRANCH_INFO 65:0

`define BRANCH_INFO_VALID 65
`define BRANCH_INFO_PC 64:33
`define BRANCH_INFO_DIR 32
`define BRANCH_INFO_TAR 31:0

// btb: direct mapped
module bpu #(
  parameter BTB_DEPTH = 256,
  parameter TAG_WIDTH = 32 - ($clog2(BTB_DEPTH) + 2),
  parameter BTB_WIDTH = TAG_WIDTH + 33,
  parameter PHT_DEPTH = 256
) (
  input clk,
  input rst,
  input [31:0] pc,
  
  // branch infomation from backend
  input [`SIZE_BRANCH_INFO] branch_info, // pc + br_taken [1bit] + br_target [32bit]
    
  // predict result
  output [32:0] pred_result,   //br_taken [1bit] + br_target [32bit]
  output [1:0] br_offset
);
  
  // update signals
  wire [31:0] corr_pc;
  wire [TAG_WIDTH-1:0] corr_btb_tag;    
  wire [$clog2(BTB_DEPTH)-1:0] corr_btb_index;                          
  wire [$clog2(PHT_DEPTH)-1:0] corr_pht_index;                     

  // lookup signals
  wire [TAG_WIDTH-1:0] btb_tag;    
  wire [$clog2(BTB_DEPTH)-1:0] btb_index;                          
  wire [$clog2(PHT_DEPTH)-1:0] pht_index;           

  // prediction results
  wire [31:0]    pred_tar;  
  wire           pred_dir;   
  wire [3:0] hit;
  wire [3:0] rstrb;
  wire [3:0] dir;

  // pht connection
  wire [$clog2(PHT_DEPTH)-1:0] pht_waddr;
  wire pht_we;
  wire pht_taken;
  wire [$clog2(PHT_DEPTH)-1:0] pht_raddr;
  wire [3:0] pht_rdata;

  // btb connection
  wire [$clog2(BTB_DEPTH)-1:0] btb_waddr;
  wire btb_we;
  wire [BTB_WIDTH-1:0] btb_wdata;
  wire [$clog2(BTB_DEPTH)-1:0] btb_raddr;
  wire [BTB_WIDTH*4-1:0] btb_rdata;
  wire [BTB_WIDTH-1:0] btb_rdata_vec [3:0];

  // btb lookup result decompose
  wire [TAG_WIDTH-1:0] tag_r [3:0];
  wire [31:0]        target_r [3:0];
  wire               valid[3:0];

  assign corr_pc = branch_info[`BRANCH_INFO_PC];
  assign corr_btb_tag = corr_pc[31:31-TAG_WIDTH+1];
  assign corr_btb_index = corr_pc[$clog2(BTB_DEPTH)+1:2];
  assign corr_pht_index = corr_pc[$clog2(PHT_DEPTH)+1:2];

  // align pht and btb read index
  assign btb_tag = pc[31:31-TAG_WIDTH+1];
  assign btb_index = {pc[$clog2(BTB_DEPTH)+1:4], 2'b00};
  assign pht_index = {pc[$clog2(PHT_DEPTH)+1:4], 2'b00};

  assign pred_result = |hit ? {pred_dir, pred_tar} : 33'd0;
  // br_offset is the pc offset of the first br_pre_taken instruction
  assign br_offset = 
    dir[0] & hit[0] ? 0 :
    dir[1] & hit[1] ? 1 :
    dir[2] & hit[2] ? 2 :
    dir[3] & hit[3] ? 3 :
                      3;
  
  // target prediction
  assign rstrb =
    {4{pc[3:2] == 0}} & 4'b1111 |
    {4{pc[3:2] == 1}} & 4'b1110 |
    {4{pc[3:2] == 2}} & 4'b1100 |
    {4{pc[3:2] == 3}} & 4'b1000 ;
  assign hit = {
    (tag_r[3] == btb_tag) && valid[3] && rstrb[3],
    (tag_r[2] == btb_tag) && valid[2] && rstrb[2],
    (tag_r[1] == btb_tag) && valid[1] && rstrb[1],
    (tag_r[0] == btb_tag) && valid[0] && rstrb[0]
  };
  assign dir = pht_rdata;
  assign pred_dir = |(hit & dir);
  assign pred_tar = 
    dir[0] & hit[0] ? target_r[0] :
    dir[1] & hit[1] ? target_r[1] :
    dir[2] & hit[2] ? target_r[2] :
    dir[3] & hit[3] ? target_r[3] :
                  32'd0;

  assign pht_waddr = corr_pht_index;
  assign pht_we = branch_info[`BRANCH_INFO_VALID];
  assign pht_taken = branch_info[`BRANCH_INFO_DIR];
  assign pht_raddr = pht_index;

  pht #(
    .DEPTH(PHT_DEPTH)
  ) pht (
    .clk(clk),
    .rst(rst),

    .waddr(pht_waddr),
    .we(pht_we),
    .taken(pht_taken),

    .raddr(pht_raddr),
    .rdata(pht_rdata)
  );

  assign btb_waddr = corr_btb_index;
  // only store the branch info of taken branches
  assign btb_we = branch_info[`BRANCH_INFO_VALID] && branch_info[`BRANCH_INFO_DIR];
  assign btb_wdata = {
    1'b1,
    corr_btb_tag,
    branch_info[`BRANCH_INFO_TAR]
  };
  assign btb_raddr = btb_index;

  genvar i;
  generate
    for (i = 0; i < 4; i = i + 1) begin
      assign btb_rdata_vec[i] = btb_rdata[(i+1)*BTB_WIDTH-1:i*BTB_WIDTH];
      assign valid[i] = btb_rdata_vec[i][BTB_WIDTH-1];
      assign tag_r[i] = btb_rdata_vec[i][BTB_WIDTH-2:32];
      assign target_r[i] = btb_rdata_vec[i][31:0];
    end
  endgenerate

  btb #(
    .DEPTH(BTB_DEPTH),
    .WIDTH(BTB_WIDTH)
  ) btb (
    .clk(clk),
    .rst(rst),

    .waddr(btb_waddr),
    .we(btb_we),
    .wdata(btb_wdata),

    .raddr(btb_raddr),
    .rdata(btb_rdata)
  );

endmodule